#!/usr/bin/python
#-------------------------#
# pymd5hashreverse.py     #
#-------------------------#
import argparse
import pyopencl as cl
import numpy as np
from math import pow

clmem = cl.mem_flags
clqueue = cl.command_queue_properties



alphabet = np.array([ord(x) for x in ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h',
                                      'i', 'j', 'k', 'l', 'm', 'n', 'o', 'p',
                                      'q', 'r', 's', 't', 'u', 'v', 'w', 'x',
                                      'y', 'z', '0', '1', '2', '3', '4', '5',
                                      '6', '7', '8', '9', 'A', 'B', 'C', 'D',
                                      'E', 'F', 'G', 'H', 'I', 'Y', 'K', 'L',
                                      'M', 'N', 'O', 'P', 'Q', 'R', 'S', 'T',
                                      'U', 'V', 'W', 'X', 'Y', 'Z', '_', '-']], dtype=np.uint8)

initStateBytes = np.array([0x67452301, 0xEFCDAB89, 0x98BADCFE, 0x10325476], dtype=np.uint32)


class MD5HashReverse:

    def __init__(self, profiling=False):
        self.profiling = profiling
        self._initCL()

    
    def _initCL(self):
        platforms = cl.get_platforms();
        if len(platforms) == 0:
            print("Failed to find any OpenCL platform")
            return None
        
        # Get avaible GPU devices and choose the first one
        # TODO: select among available gpus?
        devices = platforms[0].get_devices(cl.device_type.GPU)
        #devices = platforms[0].get_devices(cl.device_type.CPU)
        if len(devices) == 0:
            print("Could not find GPU devices")
            return None
        self.gpuDevice = devices[0]

        # Create a context using the first GPU device found
        self.context = cl.Context([self.gpuDevice])

        # Finally create the command queue
        if self.profiling:
            self.cmdQueue = cl.CommandQueue(self.context, self.gpuDevice, 
                                            properties=clqueue.PROFILING_ENABLE)
        else:
            self.cmdQueue = cl.CommandQueue(self.context, self.gpuDevice)



    def reverse(self, toReverseHash):
        kernelFile = open('md5.cl', 'r')
        kernelSrc = kernelFile.read()
        
        program = cl.Program(self.context, kernelSrc)
        program.build(devices=[self.gpuDevice])
        
	# Subtract initial final values to input hash: it results in 4 less addictions
        # performed per thread
        toReverseHash.dtype = np.uint32
        toReverseHash -= initStateBytes
        toReverseHash.dtype = np.uint8

        alphBuff = cl.Buffer(self.context, clmem.READ_ONLY|clmem.COPY_HOST_PTR,
                             hostbuf=alphabet)
        resBuff = cl.Buffer(self.context, clmem.WRITE_ONLY, size=12)
        
        inHashBuff = cl.Buffer(self.context, clmem.READ_ONLY|clmem.COPY_HOST_PTR,
                               hostbuf=toReverseHash)
        #foundBuff = cl.Buffer(self.context, clmem.WRITE_ONLY|clmem.COPY_HOST_PTR,
        #                      hostbuf=np.array([0], dtype=np.uint32))
        
        resBuff = cl.Buffer(self.context, clmem.WRITE_ONLY|clmem.COPY_HOST_PTR, 12,
                            hostbuf=np.zeros(12, dtype=np.uint8))
        result = np.empty(12, dtype=np.uint8)

        # TODO: automatically select subspace size, maybe depending on specified frequency of
        # output info
        passOffset = 0
        passSpaceSize = int(pow(2,12))
        passSubSpaceSize = int(pow(2,30))

        for passLen in range(11)[1:10]:

            if self.profiling:
                print "\nTrying passphrases of length " + str(passLen)
                print "Passphrases space: {0} items".format(passSpaceSize)

            if passSpaceSize < 256:
		ev = program.md5hash(self.cmdQueue, (256,), (256,),                
		#ev = program.md5hash(self.cmdQueue, (128,), (128,),
                                     inHashBuff, np.uint64(passOffset), alphBuff, np.uint32(passLen),
                                     #foundBuff, resBuff)
                                     resBuff)
                if self.profiling:
                    pass
                

            elif passSpaceSize < passSubSpaceSize:
                ev = program.md5hash(self.cmdQueue, (passSpaceSize,), (256,),
                #ev = program.md5hash(self.cmdQueue, (passSpaceSize,), (128,),
                                     inHashBuff, np.uint64(passOffset), alphBuff, np.uint32(passLen),
                                     #foundBuff, resBuff)
                                     resBuff)
                if self.profiling:
                    pass

                # Check if we've found the hash
                #found = np.empty(1, dtype=np.uint32)
                #cl.enqueue_copy(self.cmdQueue, found, foundBuff)
                cl.enqueue_copy(self.cmdQueue, result, resBuff)
                #if found[0]:
                if result.any():
                    #cl.enqueue_copy(self.cmdQueue, result, resBuff)
                    break

            else:
                for i in range(passSpaceSize/passSubSpaceSize):
                    ev = program.md5hash(self.cmdQueue, (passSubSpaceSize,), (256,),
                    #ev = program.md5hash(self.cmdQueue, (passSubSpaceSize,), (128,),
                                         inHashBuff, np.uint64(passOffset+(passSubSpaceSize*i)),
                                         alphBuff, np.uint32(passLen),
                                         #foundBuff, resBuff)
                                         resBuff)
                    cl.wait_for_events([ev,])
                    if self.profiling:
                        self._profile(ev, passSpaceSize, passSubSpaceSize, i+1)
                    

                    # Check if we've found the hash
                    #found = np.empty(1, dtype=np.uint32)
                    #cl.enqueue_copy(self.cmdQueue, found, foundBuff)
                    cl.enqueue_copy(self.cmdQueue, result, resBuff)
                    #if found[0]:
                    if result.any():
                        #cl.enqueue_copy(self.cmdQueue, result, resBuff)
                        break

                #if found[0]:
                if result.any():
                    break

            passOffset = 0
            passSpaceSize = int(pow(2, 6*(passLen+1)))

        return result


    def _profile(self, ev, passSpaceSize, passSubSpaceSize=0, subSpace=0):
        if passSubSpaceSize:
            percentage = (float(passSubSpaceSize)*subSpace)/passSpaceSize
            # TODO: use ev.profile.start instead of ev.profile.queued?
            execTime = (ev.profile.end - ev.profile.queued)*1e-9
            expectedExTime = execTime*((passSpaceSize/passSubSpaceSize)-subSpace)
            print "{0:.2f}% of passphrases space done in {1:.3f} seconds".format(percentage*100.0,
                                                                                 execTime)
            print "Expected remaining time: ~{0:.2f} seconds".format(expectedExTime)
        else:
            pass


    # Only for benchmarching purposes
    def _nooptimized_reverse(self, toReverseHash):
        #kernelFile = open('md5_constant.cl', 'r')
        kernelFile = open('md5_shared.cl', 'r')
        kernelSrc = kernelFile.read()
        
        program = cl.Program(self.context, kernelSrc)
        program.build(devices=[self.gpuDevice])
        
        alphBuff = cl.Buffer(self.context, clmem.READ_ONLY|clmem.COPY_HOST_PTR,
                             hostbuf=alphabet)
        
        resBuff = cl.Buffer(self.context, clmem.WRITE_ONLY, size=12)
        
        inHashBuff = cl.Buffer(self.context, clmem.READ_ONLY|clmem.COPY_HOST_PTR,
                               hostbuf=toReverseHash)
        #foundBuff = cl.Buffer(self.context, clmem.WRITE_ONLY|clmem.COPY_HOST_PTR,
        #                      hostbuf=np.array([0], dtype=np.uint32))
        
        resBuff = cl.Buffer(self.context, clmem.WRITE_ONLY|clmem.COPY_HOST_PTR, 12,
                            hostbuf=np.zeros(12, dtype=np.uint8))
        result = np.empty(12, dtype=np.uint8)

        # TODO: automatically select subspace size, maybe depending on specified frequency of
        # output info
        passOffset = 0
        passSpaceSize = int(pow(2,12))
        passSubSpaceSize = int(pow(2,30))

        for passLen in range(11)[1:10]:

            if self.profiling:
                print "\nTrying passphrases of length " + str(passLen)
                print "Passphrases space: {0} items".format(passSpaceSize)

            if passSpaceSize < 256:
		ev = program.md5hash(self.cmdQueue, (256,), (256,),                
		#ev = program.md5hash(self.cmdQueue, (128,), (128,),
                                     #inHashBuff, np.uint64(passOffset), np.uint32(passLen),
                                     inHashBuff, np.uint64(passOffset), alphBuff, np.uint32(passLen),
                                     resBuff)
                if self.profiling:
                    pass
                

            elif passSpaceSize < passSubSpaceSize:
                ev = program.md5hash(self.cmdQueue, (passSpaceSize,), (256,),
                #ev = program.md5hash(self.cmdQueue, (passSpaceSize,), (128,),
                                     #inHashBuff, np.uint64(passOffset), np.uint32(passLen),
                                     inHashBuff, np.uint64(passOffset), alphBuff, np.uint32(passLen),
                                     resBuff)
                if self.profiling:
                    pass

                # Check if we've found the hash
                #found = np.empty(1, dtype=np.uint32)
                #cl.enqueue_copy(self.cmdQueue, found, foundBuff)
                cl.enqueue_copy(self.cmdQueue, result, resBuff)
                #if found[0]:
                if result.any():
                    #cl.enqueue_copy(self.cmdQueue, result, resBuff)
                    break

            else:
                for i in range(passSpaceSize/passSubSpaceSize):
                    ev = program.md5hash(self.cmdQueue, (passSubSpaceSize,), (256,),
                    #ev = program.md5hash(self.cmdQueue, (passSubSpaceSize,), (128,),
                                         inHashBuff, np.uint64(passOffset+(passSubSpaceSize*i)),
                                         #np.uint32(passLen),
                                         alphBuff, np.uint32(passLen),
                                         resBuff)
                    cl.wait_for_events([ev,])
                    if self.profiling:
                        self._profile(ev, passSpaceSize, passSubSpaceSize, i+1)
                    

                    # Check if we've found the hash
                    #found = np.empty(1, dtype=np.uint32)
                    #cl.enqueue_copy(self.cmdQueue, found, foundBuff)
                    cl.enqueue_copy(self.cmdQueue, result, resBuff)
                    #if found[0]:
                    if result.any():
                        #cl.enqueue_copy(self.cmdQueue, result, resBuff)
                        break

                #if found[0]:
                if result.any():
                    break

            passOffset = 0
            passSpaceSize = int(pow(2, 6*(passLen+1)))

        return result




def main():
    parser = argparse.ArgumentParser(description="MD5 hash reversing using OpenCL")
    parser.add_argument('--profiling', action='store_true')
    parser.add_argument('hash')

    args = parser.parse_args()

    # TODO: perform additional check on input hash
    toRevHash = np.array([int(args.hash[x:x+2],16) for x in range(len(args.hash))[::2]],
                         dtype=np.uint8)

    reverser = MD5HashReverse(args.profiling)
    result = reverser.reverse(toRevHash)
    #result = reverser._nooptimized_reverse(toRevHash)
    
    # TODO: nicer printing of result
    if result.any():
        passphrase = ''.join([chr(x) for x in result])
        print "Hash Reversed: [{0}]".format(passphrase[:passphrase.find('\x80')])
    else:
        print ":("



if __name__=='__main__':
    exit(main())
